package cn.oldsix.spark.core.moviegroom.ml

import org.apache.spark.mllib.recommendation._
import org.apache.spark.sql.SparkSession
import org.apache.spark.storage.StorageLevel

/**
  * @ Author : Wu.D.J
  * @ Create : 2017.08.08
  */
object ModelTrain {
    def main(args: Array[String]): Unit = {
        val spark = SparkSession
                .builder()
                .appName("model-train")
                .enableHiveSupport()
                .getOrCreate()
        
        val ratingRDD = spark.sql("select * from train_data").rdd.map(x => Rating(x.getInt(0), x.getInt(1), x.getDouble(2)))
        
        ratingRDD.persist(StorageLevel.MEMORY_AND_DISK)
        
        val trainRDD = ratingRDD.map{
            case Rating(userId, movieId, rating) => (userId, movieId)
        }
        
        trainRDD.persist(StorageLevel.MEMORY_AND_DISK)
        
        val testRDD = spark.sql("select * from test_data").rdd.map(x => Rating(x.getInt(0), x.getInt(1), x.getDouble(2))).map{
            case Rating(userId, movieId, rating) => ((userId, movieId), rating)
        }
        
        testRDD.persist(StorageLevel.MEMORY_AND_DISK)
        
        val rank = 1
        
        val lambda = List(0.001, 0.005, 0.01, 0.015, 0.02, 0.1)
        
        val iteration = List(10, 20, 30, 40)
        
        var BEST_RMSE = Double.MaxValue
        
        var BEST_ITREATION = 0
        
        var BEST_LAMBDA = 0.0
        
        for (l <- lambda; i <- iteration) {
            val model = ALS.train(ratingRDD, rank, i, l)
            val predict = model.predict(trainRDD).map{
                case Rating(userId, movieId, rating) => ((userId, movieId), rating)
            }
            val predictAndFact = predict.join(testRDD)
            val MSE = predictAndFact.map{
                case ((user, product), (r1, r2)) =>
                    val err = r1 - r2
                    err * err
            }.mean()
            val RMSE = math.sqrt(MSE)
            if (RMSE < BEST_RMSE) {
                model.save(spark.sparkContext, s"/movie/model/$RMSE")
                BEST_RMSE = RMSE
                BEST_ITREATION = i
                BEST_LAMBDA = l
                println(s"Best Model is located in /movie/model/$RMSE")
                println(s"Best RMSE is $RMSE")
                println(s"Best Iteration is $BEST_ITREATION")
                println(s"Best lambda is $BEST_LAMBDA")
            }
        }
    }
}
